package com.liyunc.demo.detect.recognition;

import com.liyunc.demo.detect.utils.Constants;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.opencv_core;
import org.bytedeco.javacpp.opencv_face;

import java.io.File;
import java.nio.IntBuffer;
import java.util.LinkedList;
import java.util.List;

import static org.bytedeco.javacpp.opencv_core.CV_32SC1;
import static org.bytedeco.javacpp.opencv_imgcodecs.IMREAD_GRAYSCALE;
import static org.bytedeco.javacpp.opencv_imgcodecs.imread;
import static org.bytedeco.javacpp.opencv_imgproc.resize;

/**
 * 训练模型，生成模型文件.
 */
@Slf4j
public class TrainFromDirectory {

    /**
     * 从指定目录下训练.
     *
     * @param dirs       文件库路径
     * @param outputPath 训练模型输出路径
     */
    public void train(String[] dirs, String outputPath) {
        int total = 0;

        // 统计每个路路径下面的照片数，加在一起就是照片总数
        for (String dir : dirs) {
            List<String> files = getAllFilePath(dir);
            total += files.size();
        }
        log.info("total train file: {}", total);

        // 这里用来保存每一张照片的序号，和照片的Mat对象
        opencv_core.MatVector imageIndexMatMap = new opencv_core.MatVector(total);

        // 这里用来保存每一张照片的序号，和照片的类别
        opencv_core.Mat labels = new opencv_core.Mat(total, 1, CV_32SC1);
        IntBuffer labelsBuf = labels.createBuffer();

        int kindIndex = 1; // 类别序号，从1开始，dirs中的每个项目就是一个类别
        int imageIndex = 0; // 照片序号，从0开始

        // 每个目录下的照片都要遍历
        for (String dir : dirs) {
            // 得到当前目录下所有照片的绝对路径
            List<String> files = getAllFilePath(dir);
            // 处理一个目录下的每张照片，它们的序号不同，类别相同
            for (String file : files) {
                // imageIndexMatMap放的是照片的序号和Mat对象
                imageIndexMatMap.put(imageIndex, read(file));
                // babelsBuf放的是照片序号和类别
                labelsBuf.put(imageIndex, kindIndex);
                // 照片序号加一
                imageIndex++;
            }
            // 类别加一
            kindIndex++;
        }

        // 实例化人脸识别类
        opencv_face.FaceRecognizer faceRecognizer = opencv_face.FisherFaceRecognizer.create();
        // 训练
        faceRecognizer.train(imageIndexMatMap, labels);
        // 训练完成后，模型保存在指定位置
        faceRecognizer.save(outputPath);
        // 释放资源
        faceRecognizer.close();
    }

    /**
     * 读取指定图片的灰度图片，调整为指定大小.
     *
     * @param path 文件路径
     */
    private opencv_core.Mat read(String path) {
        opencv_core.Mat faceMat = imread(path, IMREAD_GRAYSCALE);
        resize(faceMat, faceMat, new opencv_core.Size(
                Constants.RESIZE_WIDTH, Constants.RESIZE_HEIGHT));
        return faceMat;
    }

    /**
     * 把指定路径下所有文件的绝对路径放入list集合中返回.
     *
     * @param dir 路径
     */
    private List<String> getAllFilePath(String dir) {
        List<String> paths = new LinkedList<>();
        File file = new File(dir);
        if (file.exists()) {
            // 列出该路径下的所有文件
            File[] files = file.listFiles();
            if (null != files) {
                for (File f : files) {
                    if (!f.isDirectory()) {
                        // 将每个文件的绝对路径添加到list中
                        paths.add(f.getAbsolutePath());
                    }
                }
            }
        }
        return paths;
    }
}
